/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkFlvDemuxer.h"
#include "MkMedia/MkMediaUtil.h"
#include "MkUtil/MkLog.h"

CMkFlvDemuxer::CMkFlvDemuxer()
    : m_GetFlvHeader(FALSE)
    , m_FileSize(0)
    , m_Duration(0)
{

}

Uint32 CMkFlvDemuxer::InputData(Uint8* pBuf, Uint32 nLen, Uint32& nUsed)
{
    Uint32 ErrorCode = NoneError;
    Uint32 nCurUsed = 0;
    if (!m_GetFlvHeader) {
        ErrorCode = ParseHeader(pBuf, nLen, nCurUsed);
        MkCheckErrorLog(ErrorCode, "parse flv header failed\n");
        if (nLen < nCurUsed + 4) {
            return NotHasEnoughBuffer;
        }
        m_GetFlvHeader = TRUE;
        Uint32 FirstTagSize = pBuf[nCurUsed] << 24 | pBuf[nCurUsed + 1] << 16 | pBuf[nCurUsed + 2] << 8 | pBuf[nCurUsed + 3];
        if (0 != FirstTagSize) {
            return InvalidParams;
        }
        nCurUsed += 4;
    }

    Uint32 nOffset = nCurUsed;
    while (nOffset < nLen) {
        ErrorCode = ParseTag(pBuf + nOffset, nLen - nOffset, nCurUsed);
        if (NoneError != ErrorCode) {
            break;
        }
        nOffset += nCurUsed;
    }
    nUsed = nOffset;
    return NoneError;
}

Uint32 CMkFlvDemuxer::GetEsPacket(MkEsPacket &EsPacket)
{
    if (m_EsQueue.empty()) {
        return QueueEmpty;
    }
    EsPacket = m_EsQueue.front();
    m_EsQueue.pop();
    return NoneError;
}

Uint32 CMkFlvDemuxer::ParseHeader(Uint8* pBuf, Uint32 nLen, Uint32& nUsed)
{
    if (nLen < MkFlvHeaderSize) {
        return NotHasEnoughBuffer;
    }
    Uint32 HeaderSize = pBuf[5] << 24 | pBuf[6] << 16 | pBuf[7] << 8 | pBuf[8];
    if (0 != memcmp("FLV", pBuf, 3)
        || 0x01 != pBuf[3]
        || MkFlvHeaderSize != HeaderSize) {
        return InvalidParams;
    }
    m_HasVideo = pBuf[4] & 0x01;        // bit 8
    m_HasAudio = (pBuf[4] >> 2) & 0x01; // bit 6
    nUsed = MkFlvHeaderSize;
    return NoneError;
}

Uint32 CMkFlvDemuxer::ParseTag(Uint8* pBuf, Uint32 nLen, Uint32& nUsed)
{
    if (nLen < MkFlvTagHeaderSize) {
        return NotHasEnoughBuffer;
    }
    BOOL bFilter = pBuf[0] >> 5 & 0x01;//bit 3
    //todo if 1 buffer mast decoder
    Uint8 TagType = pBuf[0] & 0x1f; //last 5 bit
    Uint32 DataSize = pBuf[1] << 16 | pBuf[2] << 8 | pBuf[3];
    Uint32 TimeStamp = pBuf[7] << 24 | pBuf[4] << 16 | pBuf[5] << 8 | pBuf[6];
    Uint32 StreamId = pBuf[8] << 16 | pBuf[9] << 8 | pBuf[10];
    if (DataSize + MkFlvTagHeaderSize + sizeof(Uint32) > nLen) {
        return NotHasEnoughBuffer;
    }
    MkEsPacket EsPacket;
    if (MkFlvTagAudio == TagType
        && m_HasAudio) {
        AudioTagParse(pBuf + MkFlvTagHeaderSize, DataSize, TimeStamp, bFilter, m_AudioParam, EsPacket);
    } else if (MkFlvTagVideo == TagType
        && m_HasVideo) {
        VideoTagParse(pBuf + MkFlvTagHeaderSize, DataSize, TimeStamp, bFilter, m_VideoParam, EsPacket);
    } else if (MkFlvTagScript == TagType) {
        ScriptTagParse(pBuf + MkFlvTagHeaderSize, DataSize, bFilter);
    }
    if (!EsPacket.BufferList.Empty()) {
        m_EsQueue.push(EsPacket);
    }
    nUsed = DataSize + MkFlvTagHeaderSize;
    Uint32 TagSize = pBuf[nUsed] << 24 | pBuf[nUsed + 1] << 16 | pBuf[nUsed + 2] << 8 | pBuf[nUsed + 3];
    if (TagSize != nUsed) {
        return InvalidParams;
    }
    nUsed += sizeof(Uint32);
    return NoneError;
}

Uint32 CMkFlvDemuxer::AudioTagParse(const Uint8 *pBuf, Uint32 nLen, Uint32 TimeStamp, BOOL Encrypted, CMkAudioParam& AudioParam, MkEsPacket& EsPacket)
{
    Uint8 Codec = pBuf[0] >> 4; //first 4 bit
    Uint8 SampleRate = pBuf[0] >> 2 & 0x03; //next 2bit
    Uint8 SampleSize = pBuf[0] >> 1 & 0x01; // next 1 bit
    Uint8 SndType = pBuf[0] & 0x01; //last 1bit
    AudioParam.SetStereo(SndType);
    AudioParam.SetSampleSize(MkAudioSampleSize(SampleSize));
    if (MkAudioRtmpAac == Codec) {
        AudioParam.SetCodec(MkCodecAac);
        if (0x00 == pBuf[1]) {
            /* AudioSpecificConfig  Iso 14496-3
                AudioObjectType 5bit
                SampleRateIndex 4bit
                ChannelConfigutation 4bit
                FrameLenghFlag 1bit
                DependOnCoreCoder 1bit
                ExternsionFlag 1bit
            */
            Uint8 AacObjectType = pBuf[2] >> 3;//first 5 bit
            SampleRate = ((pBuf[2] & 0x07) << 1) | ((pBuf[3] >> 7) & 0x01);// 2byte last 3 bit and 3byte 1bit
            Uint8 ChannelConfig = (pBuf[3] >> 3) & 0x0f; //next 4 bit
            Uint8 FrameLenghFlag = (pBuf[3] >> 2) & 0x01; //sixth bit
            Uint8 DependOnCoreCoder = (pBuf[3] >> 1) & 0x01;//seventh bit
            Uint8 ExternsionFlag = pBuf[3] & 0x01; //last bit
            AudioParam.SetObjectType((MkAacObjectType)AacObjectType);
            AudioParam.SetChannel(ChannelConfig);
            AudioParam.SetSampletRate((MkAudioSampleRate)SampleRate);
            AudioParam.SetFrameLenghFlag(FrameLenghFlag);
            AudioParam.SetDependOnCoreCoder(DependOnCoreCoder);
            AudioParam.SetExternsionFlag(ExternsionFlag);
        } else if (0x01 == pBuf[1]) {
            EsPacket.CodecType = MkCodecAac;
            EsPacket.Dts = EsPacket.Pts = TimeStamp;
            EsPacket.BufferList.Append(pBuf + 2, nLen - 2);
        } else {
            printf("invalid aac buffer type\n");
            return InvalidParams;
        }
    }
    return NoneError;
}

Uint32 CMkFlvDemuxer::VideoTagParse(const Uint8 *pBuf, Uint32 nLen, Uint32 TimeStamp, BOOL Encrypted, CMkVideoParam& VideoParam, MkEsPacket& EsPacket)
{
    if (nLen < 1) {
        return InvalidParams;
    }
    Uint8 FrameType = pBuf[0] >> 4;
    Uint8 Codec = pBuf[0] & 0x0f;
    if (MkVideoRtmpAvc == Codec) {
        VideoParam.SetCodec(MkCodecH264);
        Uint8 AvcPacketType = pBuf[1];
        Uint8 CompositionTime = pBuf[2] << 16 | pBuf[3] << 8 | pBuf[4];
        if (MkFlvAvcPacketSequenceHeader == AvcPacketType) {
            //MkH264Param Param;
            //Param.ConfigurationVersion = pBuf[5];
            //Param.Profile= pBuf[6];
            //Param.Compatibility= pBuf[7];
            //Param.Level = pBuf[8];
            Uint32 NaluLen = (pBuf[9] & 0x03) + 1;
            //printf("nalu len:%d\n", NaluLen);
            Uint16 SpsSize = pBuf[11] << 8 | pBuf[12];
            Uint16 PpsSize = pBuf[13 + SpsSize + 1] << 8 | pBuf[13 + SpsSize + 2];
            VideoParam.SetPps(MkString((Int8*)(pBuf + 13 + SpsSize + 3), PpsSize));
            VideoParam.SetSps(MkString((Int8*)(pBuf + 13), SpsSize));
            //m_VideoParam.SetH264Param(Param);
            //todo more
        } else if (MkFlvAvcPacketNalu == AvcPacketType) {
            nLen -= 5;
            const Uint8 *pOffset = pBuf + 5;
            EsPacket.CodecType = MkCodecH264;
            EsPacket.Pts = TimeStamp + CompositionTime;
            EsPacket.Dts = TimeStamp;
            while (pOffset < pBuf + nLen) {
                Uint32 NaluSize = pOffset[0] << 24 | pOffset[1] << 16 | pOffset[2] << 8 | pOffset[3];
                pOffset += sizeof(Uint32);
                EsPacket.BufferList.Append(pOffset, NaluSize);
                Uint8 NaluType = 0;
                CMkMediaUtil::Get264NaluType(*pOffset, NaluType);
                if (MkH264NalIdrSlice == NaluType) {
                    EsPacket.bKeyPacket = TRUE;
                }
                pOffset += NaluSize;
            }
        } else if (MkFlvAvcPacketEndSequence == AvcPacketType) {
            EsPacket.Pts = TimeStamp + CompositionTime;
            EsPacket.Dts = TimeStamp;
            EsPacket.CodecType = VideoParam.GetCodec();
            EsPacket.bKeyPacket = TRUE;
        } else {
            MkErrorLog("recv error avc packet\n");
            return InvalidParams;
        }
    } else if (MkVideoRtmpHevC == Codec) {
        VideoParam.SetCodec(MkCodecH265);
        if (nLen < 5) {
            return InvalidParams;
        }
        Uint8 AvcPacketType = pBuf[1];
        Uint8 CompositionTime = pBuf[2] << 16 | pBuf[3] << 8 | pBuf[4];
        if (MkFlvAvcPacketSequenceHeader == AvcPacketType) {
            //MkH265Param Param = { 0 };
            //byte 5-27 sps
            //Param.ConfigurationVersion = pBuf[5];
            //Param.GeneralProfileSpace = (pBuf[6] >> 6) & 0x03;
            //Param.GeneralTierFlag = (pBuf[6] >> 5) & 0x01;
            //Param.GeneralProfileIdc = pBuf[6] & 0x1F;
            //Param.GeneralProfileCompatibilityFlags = (pBuf[7] << 24) | (pBuf[8] << 16) | (pBuf[9] << 8) | pBuf[10];
            //Param.GeneralConstraintIndicatorFlags = ((Uint64)pBuf[11] << 24) | ((Uint64)pBuf[12] << 16) | ((Uint64)pBuf[13] << 8) | (Uint64)pBuf[14];
            //Param.GeneralConstraintIndicatorFlags = (Param.GeneralConstraintIndicatorFlags << 16) | (((Uint64)pBuf[15]) << 8) | pBuf[16];
            //Param.GeneralLevelIdc = pBuf[17];
            //Param.MinSpatialSegmentationIdc = ((pBuf[18] & 0x0F) << 8) | pBuf[19];
            //Param.ParallelismType = pBuf[20] & 0x03;
            //Param.ChromaFormat = pBuf[21] & 0x03;
            //Param.BitDepthLumaMinus8 = pBuf[22] & 0x07;
            //Param.BitDepthChromaMinus8 = pBuf[23] & 0x07;
            //Param.AvgFrameRate = (pBuf[24] << 8) | pBuf[25];
            //Param.ConstantFrameRate = (pBuf[26] >> 6) & 0x03;
            //Param.NumTemporalLayers = (pBuf[26] >> 3) & 0x07;
            //Param.TemporalIdNested = (pBuf[26] >> 2) & 0x01;
            //Param.LengthSizeMinusOne = pBuf[26] & 0x03;
            //m_VideoParam.SetH265Param(Param);
            //FILE *fp = fopen("C:/Workspace/media/tmp.dat", "wb");
            //fwrite(pBuf+5, 1, 22, fp);
            //fclose(fp);
            if (nLen < 27) {
                return InvalidParams;
            }
            Uint8 ArraySize = pBuf[27];
            const Uint8 *pOffset = pBuf + 28;
            for (Uint8 i = 0; i < ArraySize; i++) {
                if (pOffset + 3 > pBuf + nLen) {
                    return InvalidParams;
                }
                Uint8 NaluType = pOffset[0];
                Uint16 nCount = (pOffset[1] << 8) | pOffset[2];
                pOffset += 3;
                for (Uint16 j = 0; j < nCount; j++) {
                    if (pOffset + 2 > pBuf + nLen) {
                        return InvalidParams;
                    }
                    Uint16 nNaluLen = (pOffset[0] << 8) | pOffset[1];
                    if (pOffset + 2 + nNaluLen > pBuf + nLen
                        || (NaluType & 0x3F) != ((pOffset[2] >> 1) & 0x3F)) {
                        return InvalidParams;
                    }
                    pOffset += 2;
                    CMkMediaUtil::Get265NaluType(*pOffset, NaluType);
                    if (MkHevcNalVps == NaluType) {
                        VideoParam.SetVps(MkString((Int8*)pOffset, nNaluLen));
                    } else if (MkHevcNalSps == NaluType) {
                        VideoParam.SetSps(MkString((Int8*)pOffset, nNaluLen));
                    } else if (MkHevcNalPps == NaluType) {
                        VideoParam.SetPps(MkString((Int8*)pOffset, nNaluLen));
                    }
                    pOffset += nNaluLen;
                }
            }
        } else if (MkFlvAvcPacketNalu == AvcPacketType) {
            nLen -= 5;
            const Uint8 *pOffset = pBuf + 5;
            EsPacket.CodecType = MkCodecH265;
            EsPacket.Pts = TimeStamp + CompositionTime;
            EsPacket.Dts = TimeStamp;
            while (pOffset < pBuf + nLen) {
                Uint32 NaluSize = pOffset[0] << 24 | pOffset[1] << 16 | pOffset[2] << 8 | pOffset[3];
                pOffset += sizeof(Uint32);
                EsPacket.BufferList.Append(pOffset, NaluSize);
                Uint8 NaluType = 0;
                CMkMediaUtil::Get265NaluType(*pOffset, NaluType);
                if (MkHevcNalIdrWRadl == NaluType
                    || MkHevcNalIdrNLp == NaluType) {
                    EsPacket.bKeyPacket = TRUE;
                }
                pOffset += NaluSize;
            }
        } else if (MkFlvAvcPacketEndSequence == AvcPacketType) {
            EsPacket.Pts = TimeStamp + CompositionTime;
            EsPacket.Dts = TimeStamp;
            EsPacket.CodecType = VideoParam.GetCodec();
            EsPacket.bKeyPacket = TRUE;
        } else {
            MkErrorLog("recv error avc packet\n");
            return InvalidParams;
        }
    }
    return NoneError;
}

Uint32 CMkFlvDemuxer::ScriptTagParse(Uint8 *pBuf, Uint32 nLen, BOOL Encrypted)
{
    MkVector<MkAmfObject> VecObject;
    CMkAmf::DecodeObject(pBuf, pBuf + nLen, VecObject, FALSE);
    Uint32 Width = 0;
    Uint32 Height = 0;
    Uint32 VideoBitRate = 0;
    Uint32 FrameRate = 0;
    Uint8 VideoCodecType = MkVideoRtmpUnknow;
    Uint8 AudioCodeType = MkAudioRtmpUnknow;
    Uint32 AudioBitRate = 0;
    Uint32 AudioSampleRate = 0;
    Uint32 AudioSampleSize = 0;
    BOOL bStereo = FALSE;
    AmfObjectParse(VecObject, Width, Height, VideoBitRate, FrameRate, VideoCodecType, AudioCodeType, AudioBitRate, AudioSampleRate, AudioSampleSize, bStereo, m_Duration, m_FileSize);
    m_VideoParam.SetCodec(CMkMediaUtil::RtmpVideoTypeToCodec((MkRtmpVideoType)VideoCodecType));
    m_VideoParam.SetWidth(Width);
    m_VideoParam.SetHeight(Height);
    m_VideoParam.SetFrameRate(FrameRate);
    m_VideoParam.SetBitRate(VideoBitRate);
    m_AudioParam.SetCodec(CMkMediaUtil::RtmpAudioTypeToCodec((MkRtmpAudioType)AudioCodeType));
    m_AudioParam.SetBitRate(AudioBitRate);
    m_AudioParam.SetSampletRate(CMkMediaUtil::AudioSampleRateToIndex(AudioSampleRate));
    m_AudioParam.SetSampleSize(CMkMediaUtil::AudioSampleSizeToIndex(AudioSampleSize));
    m_AudioParam.SetStereo(bStereo);
    //FILE *fp = fopen("C:/Workspace/media/script.dat", "wb");
    //fwrite(pBuf, 1, nLen, fp);
    //fclose(fp);
    return NoneError;
}

Uint32 CMkFlvDemuxer::AmfObjectParse(const MkVector<MkAmfObject>& VecObject, Uint32& Width, Uint32& Height, Uint32& VideoBitRate, Uint32& FrameRate, Uint8& VideoCodecType, Uint8& AudioCodeType, Uint32& AudioBitRate, Uint32& AudioSampleRate, Uint32& AudioSampleSize, BOOL& bStereo, Uint32& Duration, Uint32& FileSize)
{
    MkVector<MkAmfObject>::const_iterator it = VecObject.begin();
    for (; it != VecObject.end(); it++) {
        AmfObjectParse(it->VecObjectValue, Width, Height, VideoBitRate, FrameRate, VideoCodecType, AudioCodeType, AudioBitRate, AudioSampleRate, AudioSampleSize, bStereo, Duration, FileSize);
        if ("duration" == it->Name) {
            Duration = static_cast<Uint32>(it->NumberValue + 0.5);
        } else if ("width" == it->Name) {
            Width = static_cast<Uint32>(it->NumberValue);
        } else if ("height" == it->Name) {
            Height = static_cast<Uint32>(it->NumberValue);
        } else if ("videodatarate" == it->Name) {
            VideoBitRate = static_cast<Uint32>(it->NumberValue + 0.5);
        } else if ("framerate" == it->Name) {
            FrameRate = static_cast<Uint32>(it->NumberValue + 0.5);
        } else if ("videocodecid" == it->Name) {
            VideoCodecType = static_cast<Uint32>(it->NumberValue);
        } else if ("audiodatarate" == it->Name) {
            AudioBitRate = static_cast<Uint32>(it->NumberValue + 0.5);
        } else if ("audiosamplerate" == it->Name) {
            AudioSampleRate = static_cast<Uint32>(it->NumberValue);
        } else if ("audiosamplesize" == it->Name) {
            AudioSampleSize = static_cast<Uint32>(it->NumberValue);
        } else if ("stereo" == it->Name) {
            bStereo = it->bValue;
        } else if ("audiocodecid" == it->Name) {
            AudioCodeType = static_cast<Uint32>(it->NumberValue);
        } else if ("filesize" == it->Name) {
            FileSize = static_cast<Uint32>(it->NumberValue);
        }
    }
    return NoneError;
}

CMkFlvDemuxer::~CMkFlvDemuxer()
{

}